{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transfer learning refers to the process of pre-training a flexible model on a large dataset and using it later on other data with little to no training. It is one of the most outstanding 🚀 achievements in Machine Learning 🧠 and has many practical applications.\n",
    "\n",
    "For time series forecasting, the technique allows you to get lightning-fast predictions ⚡ bypassing the tradeoff between accuracy and speed (more than 30 times faster than our alreadsy fast [autoARIMA](https://github.com/Nixtla/statsforecast) for a similar accuracy).\n",
    "\n",
    "This notebook shows how to generate a pre-trained model and store it in a checkpoint to make it available to forecast new time series never seen by the model. \n",
    "\n",
    "Table of Contents<br>\n",
    "1.   Installing NeuralForecast/DatasetsForecast<br>\n",
    "2.   Load M4 Data<br>\n",
    "3.   Instantiate NeuralForecast core, Fit, and save<br>\n",
    "4.   Load pre-trained model and predict on AirPassengers<br>\n",
    "5.   Evaluate Results<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can run these experiments using GPU with Google Colab.\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/Nixtla/neuralforecast/blob/main/nbs/examples/Transfer_Learning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%capture\n",
    "# !pip install git+https://github.com/Nixtla/datasetsforecast.git@main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%capture\n",
    "# !pip install neuralforecast "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from IPython.display import display, Markdown\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from datasetsforecast.m4 import M4\n",
    "from neuralforecast.core import NeuralForecast\n",
    "from neuralforecast.models import NHITS\n",
    "from neuralforecast.utils import AirPassengersDF\n",
    "from neuralforecast.losses.numpy import mae, mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example will automatically run on GPUs if available. **Make sure** cuda is available. (If you need help to put this into production send us an email or join or community, we also offer a fully hosted solution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load M4 Data\n",
    "\n",
    "The `M4` class will automatically download the complete M4 dataset and process it.\n",
    "\n",
    "It return three Dataframes: `Y_df` contains the values for the target variables, `X_df` contains exogenous calendar features and `S_df` contains static features for each time-series (none for M4). For this example we will only use `Y_df`.\n",
    "\n",
    "If you want to use your own data just replace `Y_df`. Be sure to use a long format and have a simmilar structure than our data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_df, _, _ = M4.load(directory='./', group='Monthly', cache=True)\n",
    "Y_df['ds'] = pd.to_datetime(Y_df['ds'])\n",
    "Y_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Model Train and Save\n",
    "\n",
    "Using the `NeuralForecast.fit` method you can train a set of models to your dataset. You just have to define the `input_size` and `horizon` of your model. The `input_size` is the number of historic observations (lags) that the model will use to learn to predict `h` steps in the future. Also, you can modify the hyperparameters of the model to get a better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "horizon = 12\n",
    "stacks = 3\n",
    "models = [NHITS(input_size=5 * horizon,\n",
    "                h=horizon,\n",
    "                max_steps=100,\n",
    "                stack_types = stacks*['identity'],\n",
    "                n_blocks = stacks*[1],\n",
    "                mlp_units = [[256,256] for _ in range(stacks)],\n",
    "                n_pool_kernel_size = stacks*[1],\n",
    "                batch_size = 32,\n",
    "                scaler_type='standard',\n",
    "                n_freq_downsample=[12,4,1])]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(df=Y_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save model with `core.NeuralForecast.save` method. This method uses PytorchLightning `save_checkpoint` function. We set `save_dataset=False` to only save the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nf.save(path='./results/transfer/', model_index=None, overwrite=True, save_dataset=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Transfer M4 to AirPassengers\n",
    "\n",
    "We load the stored model with the `core.NeuralForecast.load` method, and forecast `AirPassenger` with the `core.NeuralForecast.predict` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fcst2 = NeuralForecast.load(path='./results/transfer/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We define the train df. \n",
    "Y_df = AirPassengersDF.copy()\n",
    "mean = Y_df[Y_df.ds<='1959-12-31']['y'].mean()\n",
    "std = Y_df[Y_df.ds<='1959-12-31']['y'].std()\n",
    "\n",
    "Y_train_df = Y_df[Y_df.ds<='1959-12-31'] # 132 train\n",
    "Y_test_df = Y_df[Y_df.ds>'1959-12-31']   # 12 test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_hat_df = fcst2.predict(df=Y_train_df).reset_index()\n",
    "Y_hat_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "Y_hat_df = Y_test_df.merge(Y_hat_df, how='left', on=['unique_id', 'ds'])\n",
    "plot_df = pd.concat([Y_train_df, Y_hat_df]).set_index('ds')\n",
    "\n",
    "plot_df[['y', 'NHITS']].plot(ax=ax, linewidth=2)\n",
    "\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Timestamp [t]', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate Results\n",
    "\n",
    "\n",
    "We evaluate the forecasts of the pre-trained model with the Mean Absolute Error (`mae`).\n",
    "\n",
    "$$\n",
    "\\qquad MAE = \\frac{1}{Horizon} \\sum_{\\tau} |y_{\\tau} - \\hat{y}_{\\tau}|\\qquad\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true = Y_test_df.y.values\n",
    "y_hat = Y_hat_df['NHITS'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('NHITS     MAE: %0.3f' % mae(y_hat, y_true))\n",
    "print('ETS       MAE: 16.222')\n",
    "print('AutoARIMA MAE: 18.551')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
